
open Class_ast
open Operator_ast_fullfill
open Operator_ast_absolute
open Tools_identifier

module StrSet = Set.Make(String)
module StrMap = Map.Make(String)

type context = {
	cop_operators : string StrMap.t;
	cop_idents : StrSet.t;
}

let transform ast =

	let rec transform_list func (data:context) = function
		| [] -> data, []
		| a::lst -> let data,a = func (data:context) a in
			let data,lst = transform_list func (data:context) lst in
			data,(a::lst) in
	
	let is_unary = function
			| "++" | "--" | "!" | "~" -> true
			| _ -> false in
	
	let get_func_expr context op =
		ACondExpr(ACondBinaryExpr(AUnaryExpr(APostfixExpr(ALeftHandSide(ANewExpression(
		AMemberExpression(AFunctionExpr(true,None,(if is_unary op then ["a"] else ["a";"b"]),[AStatement (AReturnStatement (Some (
			[ACondExpr(ACondBinaryExpr(if is_unary op then
				AUnaryExpr(AOpUnaryExpr(op,APostfixExpr(ALeftHandSide(
				ANewExpression(AMemberExpression(APrimaryExpr(AIdentifier "a")))))))
				else
				ABinaryExpr(op,AUnaryExpr(APostfixExpr(ALeftHandSide(
				ANewExpression(AMemberExpression(APrimaryExpr(AIdentifier "a")))))), AUnaryExpr(APostfixExpr(ALeftHandSide(
				ANewExpression(AMemberExpression(APrimaryExpr(AIdentifier "b")))))))
			))]
		)))]
		)))))))) in

	let rec transform_program ast = 
		let fullfill = fullfill_list fullfill_sourceElement {
			idents = StrSet.empty;
			operators = StrSet.empty
		} ast in
		let absolute = absolute_list absolute_sourceElement {
			abs_idents = StrSet.empty
		} fullfill.idents ast in
		
		let context = StrSet.fold (fun op context -> 
			let map,_,_ = free_ident (StrMap.empty,context.cop_idents,absolute.abs_idents) "?" false in
			let ident = StrMap.find "?" map in
			{
				cop_operators = StrMap.add op ident context.cop_operators;
				cop_idents = StrSet.add ident context.cop_idents
			} ) fullfill.operators {
			cop_operators = StrMap.empty;
			cop_idents = fullfill.idents;
		} in
		
		let _, ast = transform_list transform_sourceElement context ast in
		
		StrSet.fold (fun op ast -> 
			let v = StrMap.find op context.cop_operators in
			(AStatement (
			AVariableStatement (true, false, [
				(v, if op.[0] = '@' then None else Some (get_func_expr context op))
				])))::ast
			) fullfill.operators ast
	
	and transform_sourceElement (data:context) = function
		| Operator_ast.AStatement stmt -> 
			let data, stmt = transform_statement (data:context) stmt in
			data, (AStatement stmt)
		(*| Operator_ast.AFunctionDeclaration funcDecl ->  
			let data= funcDecl = transform_functionDeclaration (data:context) funcDecl in
			data, (AFunctionDeclaration funcDecl) *) (* Modif #1 *)

	and transform_functionDeclaration (data:context) (is_private, is_meta, is_inline,ident,params,body) = 
		let data, params, body = transform_functionContent (data:context) params body in
		data, (is_private, is_meta, is_inline, ident, params, body)
		
	and transform_functionExpression (data:context) (is_inline, ident, params, body) =
		let data, params, body = transform_functionContent (data:context) params body in
		data, (is_inline, ident, params, body)
	
	and transform_functionContent (data:context) params body =
		let fullfill = fullfill_list fullfill_sourceElement {
			idents = List.fold_right (fun i set -> StrSet.add i set ) params data.cop_idents;
			operators = StrSet.empty
		} body in
		let absolute = absolute_list absolute_sourceElement {
			abs_idents = StrSet.empty
		} fullfill.idents body in
		
		let data2 = StrSet.fold (fun op context -> 
			let map,_,_ = free_ident (StrMap.empty,context.cop_idents,absolute.abs_idents) "?" false in
			let ident = StrMap.find "?" map in
			{
				cop_operators = StrMap.add op ident context.cop_operators;
				cop_idents = StrSet.add ident context.cop_idents
			} ) fullfill.operators {
			cop_operators = data.cop_operators;
			cop_idents = fullfill.idents;
		} in
		let _, body = transform_functionBody (data2:context) body in
		
		let body = StrSet.fold (fun op ast -> 
			let v = StrMap.find op data2.cop_operators in
			(AStatement (
			AVariableStatement (true, false, [
				(v, if StrMap.mem op data.cop_operators then 
				Some (ACondExpr(ACondBinaryExpr(AUnaryExpr(APostfixExpr(ALeftHandSide(ANewExpression(
				AMemberExpression(APrimaryExpr(AIdentifier (StrMap.find op data.cop_operators)))
				)))))))
				else if op.[0] = '@' then None else Some (get_func_expr data2 op))
				])))::ast
			) fullfill.operators body in
		
		data, params, body
		
	and transform_functionBody (data:context) body = 
		transform_list transform_sourceElement (data:context) body

	(* Statements *)
	and transform_block (data:context) stmts =
		transform_list transform_statement (data:context) stmts

	and transform_statement (data:context) = function
		| Operator_ast.ABlock block -> 
			let data, block = transform_block (data:context) block in
			data, (ABlock block)
		| Operator_ast.AOperatorAssignment stmt -> transform_operatorAssignment data stmt
		| Operator_ast.AOperatorUnassignment stmt -> transform_operatorUnassignment data stmt
		| Operator_ast.AVariableStatement vs -> 
			let data, vs = transform_variableStatement (data:context) vs in
			data, (AVariableStatement vs)
		| Operator_ast.AConstStatement cs -> 
			let data, cs = transform_constStatement (data:context) cs in
			data, (AConstStatement cs)
		| Operator_ast.AFuncDeclaration fd -> 
			let data, fd = transform_functionDeclaration (data:context) fd in
			data, (AFuncDeclaration fd)	(* Modif #1 *)
		| Operator_ast.AClassStatement cs -> 
			let data, cs = transform_classStatement (data:context) cs in
			data, (AClassStatement cs)
		| Operator_ast.AEmptyStatement -> data, AEmptyStatement
		| Operator_ast.AExpressionStatement expr -> 
			let data, expr = transform_expression (data:context) expr in
			data, (AExpressionStatement expr)
		| Operator_ast.AIfStatement stmt -> 
			let data, stmt = transform_ifStatement (data:context) stmt in
			data, (AIfStatement stmt)
		| Operator_ast.AIterationStatement stmt -> 
			let data, stmt = transform_iterationStatement (data:context) stmt in
			data, (AIterationStatement stmt)
		| Operator_ast.AContinueStatement e -> data, AContinueStatement e
		| Operator_ast.ABreakStatement e -> data, ABreakStatement e
		| Operator_ast.AReturnStatement expr -> 
			let data, expr = transform_optionalExpression (data:context) expr in
			data, (AReturnStatement expr)
		| Operator_ast.AWithStatement stmt -> 
			let data, stmt = transform_withStatement (data:context) stmt in
			data, (AWithStatement stmt)
		| Operator_ast.ALabelledStatement stmt -> 
			let data, stmt = transform_labelledStatement (data:context) stmt in
			data, (ALabelledStatement stmt)
		| Operator_ast.ASwitchStatement stmt -> 
			let data, stmt = transform_switchStatement (data:context) stmt in
			data, (ASwitchStatement stmt)
		| Operator_ast.AThrowStatement expr -> 
			let data, expr = transform_expression (data:context) expr in
			data, (AThrowStatement expr)
		| Operator_ast.ATryStatement stmt -> 
			let data, stmt = transform_tryStatement (data:context) stmt in
			data, (ATryStatement stmt)
		| Operator_ast.ADebuggerStatement -> data, ADebuggerStatement
	
	and transform_operatorAssignment (data:context) (op,expr) =
		let v = StrMap.find op data.cop_operators in
		let data, expr = transform_assignmentExpression data expr in
		data, (AExpressionStatement [AAssignOp (
			ANewExpression(AMemberExpression(APrimaryExpr(AIdentifier v))),"=",
			expr
		)])
	
	and transform_operatorUnassignment (data:context) op =
		let v = StrMap.find op data.cop_operators in
		if op.[0] = '@' then data, AEmptyStatement else
		data, (AExpressionStatement [AAssignOp (
			ANewExpression(AMemberExpression(APrimaryExpr(AIdentifier v))),"=",
			(get_func_expr data op)
		)])

	and transform_variableStatement (data:context) (is_private, is_meta, varDecls) =
		let data, varDecls = transform_list transform_variableDeclaration (data:context) varDecls in
		data, (is_private, is_meta, varDecls)
		
	and transform_classStatement (data:context) (is_private, is_meta, ident, extends, classElements) =
		let data, classElements = transform_list transform_classElement (data:context) classElements in
		data, (is_private, is_meta, ident, extends, classElements)

	and transform_classElement (data:context) = function
		| Operator_ast.ACVariableStatement v -> 
			let data, v = transform_variableStatement (data:context) v in
			data, (ACVariableStatement v)
		| Operator_ast.ACConstStatement c -> 
			let data, c = transform_constStatement (data:context) c in
			data, (ACConstStatement c)
		| Operator_ast.ACFuncDeclaration funcDecl -> 
			let data, funcDecl = transform_functionDeclaration (data:context) funcDecl in
			data, (ACFuncDeclaration funcDecl)

	and transform_constStatement (data:context) (is_private, is_meta, varDecls) =
		let data, varDecls = transform_list transform_variableDeclaration (data:context) varDecls in
		data, (is_private, is_meta, varDecls)

	and transform_variableDeclaration (data:context) (ident, init) =
		(match init with
			| None -> data, (ident, None)
			| Some expr -> 
				let data, expr = transform_assignmentExpression (data:context) expr in
				data, (ident, (Some expr)))

	and transform_ifStatement (data:context) = function
		| Operator_ast.AIfElse (expr, stmt1, stmt2) -> 
			let data, expr = transform_expression (data:context) expr in
			let data, stmt1 = transform_statement (data:context) stmt1 in
			let data, stmt2 = transform_statement (data:context) stmt2 in
			data, (AIfElse (expr, stmt1, stmt2))
		| Operator_ast.AIf (expr, stmt) -> 
			let data, expr = transform_expression (data:context) expr in
			let data, stmt = transform_statement (data:context) stmt in
			data, (AIf (expr, stmt))

	and transform_iterationStatement (data:context) = function
		| Operator_ast.ADoWhile (stmt, expr) ->
			let data, stmt = transform_statement (data:context) stmt in
			let data, expr = transform_expression (data:context) expr in
			data, (ADoWhile (stmt, expr))
		| Operator_ast.AWhile (expr, stmt) -> 
			let data, expr = transform_expression (data:context) expr in
			let data, stmt = transform_statement (data:context) stmt in
			data, (AWhile (expr, stmt))
		| Operator_ast.AFor (expr1, expr2, expr3, stmt) ->
			let data, expr1 = transform_optionalExpression (data:context) expr1 in
			let data, expr2 = transform_optionalExpression (data:context) expr2 in
			let data, expr3 = transform_optionalExpression (data:context) expr3 in
			let data, stmt = transform_statement (data:context) stmt in
			data, (AFor (expr1, expr2, expr3, stmt ))
		| Operator_ast.AForVar (is_meta, varDecls, expr1, expr2, stmt) ->
			let data, varDecls = transform_list transform_variableDeclaration (data:context) varDecls in
			let data, expr1 = transform_optionalExpression data expr1 in
			let data, expr2 = transform_optionalExpression data expr2 in
			let data, stmt = transform_statement data stmt in
			data, (AForVar (is_meta, varDecls, expr1, expr2, stmt))
		| Operator_ast.AForIn (lhs, expr, stmt) ->
			let data, lhs = transform_leftHandSideExpression (data:context) lhs in
			let data, expr = transform_expression (data:context) expr in
			let data, stmt = transform_statement (data:context) stmt in
			data, (AForIn (lhs, expr, stmt))
		| Operator_ast.AForVarIn (is_meta, varDecl, expr, stmt) ->
			let data, varDecl = transform_variableDeclaration (data:context) varDecl in
			let data, expr = transform_expression (data:context) expr in
			let data, stmt = transform_statement (data:context) stmt in
			data, (AForVarIn (is_meta, varDecl, expr, stmt))

	and transform_optionalExpression (data:context) = function
		| None -> data, None
		| Some expr -> let data, expr = transform_expression (data:context) expr in
			data, (Some expr)

	and transform_withStatement (data:context) (expr, stmt) =
		let data, expr = transform_expression (data:context) expr in
		let data, stmt = transform_statement (data:context) stmt in
		data, (expr,stmt)
		
	and transform_switchStatement (data:context) (expr, caseBlock) =
		let data, expr = transform_expression (data:context) expr in
		let data, caseBlock = transform_caseBlock (data:context) caseBlock in
		data, (expr,caseBlock)

	and transform_caseBlock (data:context) = function
		| Operator_ast.ACasesNoDefault caseClauses -> 
			let data, caseClauses = transform_list transform_caseClause (data:context) caseClauses in
			data, (ACasesNoDefault caseClauses)
		| Operator_ast.ACasesDefault (caseClauses1, defaultClause, caseClauses2) ->
			let data, caseClauses1 = transform_list transform_caseClause (data:context) caseClauses1 in
			let data, defaultClause = transform_list transform_statement (data:context) defaultClause in
			let data, caseClauses2 = transform_list transform_caseClause (data:context) caseClauses2 in
			data, (ACasesDefault (caseClauses1, defaultClause, caseClauses2))

	and transform_caseClause (data:context) (expr, stmts) = 
		let data, expr = transform_expression (data:context) expr in
		let data, stmts = transform_list transform_statement (data:context) stmts in
		data, (expr, stmts)

	and transform_labelledStatement (data:context) (ident, stmt) =
		let data, stmt = transform_statement (data:context) stmt in
		data, (ident, stmt)

	and transform_tryStatement (data:context) = function
		| Operator_ast.ATryCatch (block, catch) ->
			let data, block = transform_block (data:context) block in
			let data, catch = transform_catch (data:context) catch in
			data, (ATryCatch (block, catch))
		| Operator_ast.ATryFinally (block, finally) ->
			let data, block = transform_block (data:context) block in
			let data, finally = transform_block (data:context) finally in
			data, (ATryFinally (block, finally))
		| Operator_ast.ATryCatchFinally (block, catch, finally) ->
			let data, block = transform_block (data:context) block in
			let data, catch = transform_catch (data:context) catch in
			let data, finally = transform_block (data:context) finally in
			data, (ATryCatchFinally (block, catch, finally))
			
	and transform_catch (data:context) (ident, block) =
		let data, block = transform_block (data:context) block in
		data, (ident, block)

	(* Literals *)
	and transform_literal = function
		| Operator_ast.ANull -> ANull
		| Operator_ast.ABooleanLiteral b -> ABooleanLiteral b
		| Operator_ast.ALitNumericLiteral s -> ALitNumericLiteral s
		| Operator_ast.ALitStringLiteral s -> ALitStringLiteral s
		| Operator_ast.ARegularExpressionLiteral regex -> ARegularExpressionLiteral regex

	(* Expressions *)
	and transform_primaryExpression (data:context) = function
		| Operator_ast.AThis -> data, AThis
		| Operator_ast.AIdentifier ident -> data, (AIdentifier ident)
		| Operator_ast.ALiteral literal -> data, ALiteral (transform_literal literal)
		| Operator_ast.AArrayLiteral literal -> 
			let data, literal = transform_arrayLiteral (data:context) literal in
			data, (AArrayLiteral literal)
		| Operator_ast.AObjectLiteral literal -> 
			let data, literal = transform_objectLiteral (data:context) literal in
			data, (AObjectLiteral literal)
		| Operator_ast.AExpression expr -> 
			let data, expr = transform_expression (data:context) expr in
			data, (AExpression expr)

	and transform_arrayLiteral (data:context) = function
		| Operator_ast.AElisions elisions -> data, (AElisions (List.map transform_elision elisions))
		| Operator_ast.AElements elements -> 
			let data, elements = transform_list transform_element (data:context) elements in
			data, (AElements elements)
		| Operator_ast.AElementsAndElisions (elements, elisions) ->
			let data, elements = transform_list transform_element (data:context) elements in
			let elisions = List.map transform_elision elisions in
			data, (AElementsAndElisions (elements, elisions))

	and transform_elision _ = AElision

	and transform_element (data:context) (elisions, expr) = 
		let elisions = List.map transform_elision elisions in
		let data, expr = transform_assignmentExpression data expr in
		data, (elisions,expr)

	and transform_objectLiteral (data:context) props = 
		transform_list transform_propertyAssignment (data:context) props

	and transform_propertyAssignment (data:context) = function
		| Operator_ast.APropAssignment (prop,expr) -> 
			let prop = transform_propertyName prop in
			let data, expr = transform_assignmentExpression (data:context) expr in
			data, (APropAssignment (prop, expr))
		| Operator_ast.APropGet (prop, body) ->
			let prop = transform_propertyName prop in
			let data, _, body = transform_functionContent (data:context) [] body in
			data, (APropGet (prop, body))
		| Operator_ast.APropSet (prop, params, body) ->
			let prop = transform_propertyName prop in
			let data, params, body = transform_functionContent (data:context) params body in
			data, (APropSet (prop, params, body))

	and transform_propertyName = function
		| Operator_ast.AIdentifierName s -> AIdentifierName s
		| Operator_ast.AStringLiteral s -> AStringLiteral s
		| Operator_ast.ANumericLiteral s -> ANumericLiteral s

	and transform_memberExpression (data:context) = function
		| Operator_ast.APrimaryExpr expr -> 
			let data, expr = transform_primaryExpression (data:context) expr in
			data, (APrimaryExpr expr)
		| Operator_ast.AFunctionExpr expr -> 
			let data, expr = transform_functionExpression (data:context) expr in
			data, (AFunctionExpr expr)
		| Operator_ast.AMemberBracketExpr (expr1, expr2) -> 
			let data, expr1 = transform_memberExpression (data:context) expr1 in
			let data, expr2 = transform_expression (data:context) expr2 in
			data, (AMemberBracketExpr (expr1, expr2))
		| Operator_ast.AMemberDotExpr (expr, ident) ->
			let data, expr = transform_memberExpression (data:context) expr in
			data, (AMemberDotExpr (expr, ident))
		| Operator_ast.AMemberNewExpr (expr, args) ->
			let data, expr = transform_memberExpression (data:context) expr in
			let data, args = transform_arguments (data:context) args in
			data, (AMemberNewExpr (expr, args))

	and transform_newExpression (data:context) = function
		| Operator_ast.AMemberExpression expr ->
			let data, expr = transform_memberExpression (data:context) expr in
			data, (AMemberExpression expr)
		| Operator_ast.ANewMemberExpression expr ->
			let data, expr = transform_newExpression (data:context) expr in
			data, (ANewMemberExpression expr)

	and transform_callExpression (data:context) = function
		| Operator_ast.ACallMemberArgs (expr,args) ->
			let data, expr = transform_memberExpression (data:context) expr in
			let data, args = transform_arguments (data:context) args in
			data, (ACallMemberArgs (expr, args))
		| Operator_ast.ACallExpressionArgs (expr, args) ->
			let data, expr = transform_callExpression (data:context) expr in
			let data, args = transform_arguments (data:context) args in
			data, (ACallExpressionArgs (expr, args))
		| Operator_ast.ACallExpressionExpr (expr1, expr2) ->
			let data, expr1 = transform_callExpression (data:context) expr1 in
			let data, expr2 = transform_expression (data:context) expr2 in
			data, (ACallExpressionExpr (expr1, expr2))
		| Operator_ast.ACallExpressionIdent (expr, ident) ->
			let data, expr = transform_callExpression (data:context) expr in
			data, (ACallExpressionIdent (expr, ident))

	and transform_arguments (data:context) args = 
		transform_list transform_assignmentExpression (data:context) args

	and transform_leftHandSideExpression (data:context) = function
		| Operator_ast.ANewExpression expr ->
			let data, expr = transform_newExpression (data:context) expr in
			data, (ANewExpression expr)
		| Operator_ast.ACallExpression expr ->
			let data, expr = transform_callExpression (data:context) expr in
			data, (ACallExpression expr)

	and transform_postfixExpression (data:context) = function
		| Operator_ast.ALeftHandSide lhs ->
			let data, lhs = transform_leftHandSideExpression (data:context) lhs in
			data, (ALeftHandSide lhs)
		| Operator_ast.ALeftHandSideOp (lhs,op) ->
			let data, lhs = transform_leftHandSideExpression (data:context) lhs in
			if (StrMap.mem op data.cop_operators) && op.[0] = '@' then
			let v = StrMap.find op data.cop_operators in
			data, (ALeftHandSide(ACallExpression(
			ACallMemberArgs(APrimaryExpr(AIdentifier v),
			[ACondExpr(ACondBinaryExpr (AUnaryExpr (APostfixExpr (ALeftHandSide lhs))))])
			)))
			else
			data, ALeftHandSideOp (lhs,op)

	and transform_unaryExpression (data:context) = function
		| Operator_ast.APostfixExpr expr ->
			let data, expr = transform_postfixExpression (data:context) expr in
			data, (APostfixExpr expr)
		| Operator_ast.AOpUnaryExpr (op,expr) ->
			let data, expr = transform_unaryExpression (data:context) expr in
			if (StrMap.mem op data.cop_operators) && (op.[0] = '@' || is_unary op) then
			let v = StrMap.find op data.cop_operators in
			data, (APostfixExpr(ALeftHandSide(ACallExpression(
			ACallMemberArgs(APrimaryExpr(AIdentifier v),
			[ACondExpr(ACondBinaryExpr (AUnaryExpr expr))])
			))))
			else
			data, (AOpUnaryExpr (op, expr))

	and transform_binaryExpression (data:context) = function
		| Operator_ast.AUnaryExpr expr -> 
			let data, expr = transform_unaryExpression (data:context) expr in
			data, (AUnaryExpr expr)
		| Operator_ast.ABinaryExpr (op, expr1, expr2) ->
			let data, expr1 = transform_binaryExpression (data:context) expr1 in
			let data, expr2 = transform_binaryExpression (data:context) expr2 in
			if StrMap.mem op data.cop_operators then
			let v = StrMap.find op data.cop_operators in
			data, (AUnaryExpr(APostfixExpr(ALeftHandSide(ACallExpression(
			ACallMemberArgs(APrimaryExpr(AIdentifier v),
			[ACondExpr(ACondBinaryExpr expr1);ACondExpr(ACondBinaryExpr expr2)])
			)))))
			else
			data, (ABinaryExpr (op, expr1, expr2))

	and transform_conditionalExpression (data:context) = function
		| Operator_ast.ACondBinaryExpr expr ->
			let data, expr = transform_binaryExpression (data:context) expr in
			data, (ACondBinaryExpr expr)
		| Operator_ast.ACondTernaryExpr (expr1, expr2, expr3) ->
			let data, expr1 = transform_binaryExpression (data:context) expr1 in
			let data, expr2 = transform_assignmentExpression (data:context) expr2 in
			let data, expr3 = transform_assignmentExpression (data:context) expr3 in
			data, (ACondTernaryExpr (expr1, expr2, expr3))

	and transform_assignmentExpression (data:context) = function
		| Operator_ast.ACondExpr expr -> 
			let data, expr = transform_conditionalExpression (data:context) expr in
			data, (ACondExpr expr)
		| Operator_ast.AAssignOp (lhs,op,expr) ->
			let data, lhs = transform_leftHandSideExpression (data:context) lhs in
			let data, expr = transform_assignmentExpression (data:context) expr in
			data, (AAssignOp (lhs,op,expr))

	and transform_expression (data:context) exprs = 
		transform_list transform_assignmentExpression (data:context) exprs

	in
	transform_program ast

